{
 "cells": [
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-25T01:55:39.416666Z",
     "start_time": "2025-08-25T01:55:38.108746Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "from torch.optim import AdamW\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "# 训练数据\n",
    "dataset = datasets.MNIST(root='./data/', train=True, download=True, transform=transforms.ToTensor())\n",
    "subset_indices = list(range(1000))\n",
    "dataloader = DataLoader(dataset, batch_size=64, sampler=torch.utils.data.SubsetRandomSampler(subset_indices))\n",
    "\n",
    "# 训练参数\n",
    "T = 1000\n",
    "betas = torch.linspace(0.0001, 0.02, T)\n",
    "alphas = 1 - betas\n",
    "alphas_bar = torch.cumprod(alphas, dim=0)\n",
    "sqrt_alphas_bar = torch.sqrt(alphas_bar)\n",
    "sqrt_one_minus_alphas_bar = torch.sqrt(1 - alphas_bar)"
   ],
   "id": "a3fdea946f76c3f1",
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-26T02:41:19.515199Z",
     "start_time": "2025-08-26T02:41:19.511450Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import os\n",
    "# 设置环境变量\n",
    "os.environ[\"HF_ENDPOINT\"] = \"https://hf-mirror.com\"\n",
    "os.environ[\"HF_HUB_ENABLE_HF_TRANSFER\"] = \"1\"  # 启用高速下载"
   ],
   "id": "a00ca8f6f3619ad3",
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-26T02:41:35.844273Z",
     "start_time": "2025-08-26T02:41:20.957027Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from transformers import ChineseCLIPModel, ChineseCLIPProcessor, AutoTokenizer\n",
    "\n",
    "clip_model_name = \"OFA-Sys/chinese-clip-vit-base-patch16\"\n",
    "clip_model = ChineseCLIPModel.from_pretrained(clip_model_name)\n",
    "clip_processor = ChineseCLIPProcessor.from_pretrained(clip_model_name)\n",
    "clip_tokenizer = AutoTokenizer.from_pretrained(clip_model_name)\n",
    "\n",
    "NUM_TO_TEXT = {\n",
    "    0: \"手写数字0\",\n",
    "    1: \"手写数字1\",\n",
    "    2: \"手写数字2\",\n",
    "    3: \"手写数字3\",\n",
    "    4: \"手写数字4\",\n",
    "    5: \"手写数字5\",\n",
    "    6: \"手写数字6\",\n",
    "    7: \"手写数字7\",\n",
    "    8: \"手写数字8\",\n",
    "    9: \"手写数字9\"\n",
    "}"
   ],
   "id": "af5b2f15dc62080d",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/dadudu/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. `use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`.\n"
     ]
    }
   ],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-25T01:55:53.181688Z",
     "start_time": "2025-08-25T01:55:53.171532Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torch.nn import MultiheadAttention\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class CrossAttention(nn.Module):\n",
    "    def __init__(self, dim, heads=4):\n",
    "        super().__init__()\n",
    "        self.attn = MultiheadAttention(embed_dim=dim, num_heads=heads, batch_first=True)\n",
    "        self.proj = nn.Linear(dim, dim)\n",
    "        self.norm = nn.LayerNorm(dim)\n",
    "\n",
    "    def forward(self, x, context):\n",
    "        batch_size, channel, height, width = x.shape\n",
    "        x = x.view(batch_size, channel, -1).permute(0, 2, 1)  # [batch_size, height*width, channel]\n",
    "        context = context.unsqueeze(1) # [batch_size, context_dim] -> [batch_size, 1, context_dim]\n",
    "        out, _ = self.attn(x, context, context)  # [batch_size, height*width, channel]\n",
    "        out = self.norm(out + x) # Add & Norm\n",
    "        out = self.proj(out)\n",
    "        out = out.permute(0, 2, 1).view(batch_size, channel, height, width)\n",
    "        return out\n",
    "\n",
    "# 两次卷积操作，卷积核为3，填充为1，使得输入和输出的尺寸一致\n",
    "class DoubleConv(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(DoubleConv, self).__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv2d(in_ch, out_ch, 3, padding=1),\n",
    "            nn.BatchNorm2d(out_ch),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(out_ch, out_ch, 3, padding=1),\n",
    "            nn.BatchNorm2d(out_ch),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "# 下采样，图片缩小一半\n",
    "class Down(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(Down, self).__init__()\n",
    "        self.mpconv = nn.Sequential(\n",
    "            nn.MaxPool2d(2),\n",
    "            DoubleConv(in_ch, out_ch)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.mpconv(x)\n",
    "\n",
    "# 上采样，图片放大两倍\n",
    "class Up(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(Up, self).__init__()\n",
    "        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
    "        self.conv = DoubleConv(in_ch, out_ch)\n",
    "\n",
    "    def forward(self, x1, x2):\n",
    "        x1 = self.up(x1)\n",
    "        # 裁剪并拼接跳跃连接\n",
    "        diffY = x2.size()[2] - x1.size()[2]\n",
    "        diffX = x2.size()[3] - x1.size()[3]\n",
    "        pad_left = diffX // 2\n",
    "        pad_right = diffX - pad_left\n",
    "        pad_top = diffY // 2\n",
    "        pad_bottom = diffY - pad_top\n",
    "        x1 = nn.functional.pad(x1, [pad_left, pad_right, pad_top, pad_bottom])\n",
    "        x = torch.cat([x2, x1], dim=1)\n",
    "\n",
    "        return self.conv(x)\n",
    "\n",
    "class OutConv(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):\n",
    "        super(OutConv, self).__init__()\n",
    "        # 卷积核大小为1，使得输出和输入尺寸一致，但还是变换了图片通道数\n",
    "        # out_ch=n_classes，n_classes=1，所以就是把之前的图片通道数变为1\n",
    "        # 最终整个UNet输入一张图片，输出也是一张图片\n",
    "        self.conv = nn.Conv2d(in_ch, out_ch, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # sigmoid是把像素值变为0-1之间\n",
    "        # return torch.sigmoid(self.conv(x))\n",
    "        return self.conv(x)\n",
    "\n",
    "# 大都督周瑜（我的微信: it_zhouyu）\n",
    "\n",
    "class ZhouyuUNet(nn.Module):\n",
    "    def __init__(self, n_channels=1, n_classes=1, T=1000, time_emb_dim=64):\n",
    "        super(ZhouyuUNet, self).__init__()\n",
    "\n",
    "        self.time_emb = nn.Embedding(T, time_emb_dim)\n",
    "        self.time_proj = nn.Linear(time_emb_dim, 512)\n",
    "\n",
    "        self.cross_attn = CrossAttention(dim=512)\n",
    "\n",
    "        self.inc = DoubleConv(n_channels, 64)\n",
    "        self.down1 = Down(64, 128)\n",
    "        self.down2 = Down(128, 256)\n",
    "        self.down3 = Down(256, 512)\n",
    "        self.down4 = Down(512, 512)\n",
    "        self.up1 = Up(512 + 512, 256)   # 通道数：512 (来自下采样) + 512 (上采样输入)\n",
    "        self.up2 = Up(256 + 256, 128)\n",
    "        self.up3 = Up(128 + 128, 64)\n",
    "        self.up4 = Up(64 + 64, 64)\n",
    "        self.outc = OutConv(64, n_classes)\n",
    "\n",
    "    def forward(self, x, t, text_emb=None):\n",
    "        # 时间嵌入\n",
    "        time_emb = self.time_emb(t)\n",
    "        time_emb = self.time_proj(time_emb)\n",
    "        time_emb = time_emb.view(-1, 512, 1, 1)\n",
    "\n",
    "        x1 = self.inc(x)\n",
    "        x2 = self.down1(x1)\n",
    "        x3 = self.down2(x2)\n",
    "        x4 = self.down3(x3)\n",
    "        x5 = self.down4(x4)\n",
    "        x5 = x5 + time_emb\n",
    "        x5 = self.cross_attn(x5, text_emb)\n",
    "        x = self.up1(x5, x4)\n",
    "        x = self.up2(x, x3)\n",
    "        x = self.up3(x, x2)\n",
    "        x = self.up4(x, x1)\n",
    "        x = self.outc(x)\n",
    "        return x"
   ],
   "id": "d9169d9e0955c270",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-25T01:55:53.242368Z",
     "start_time": "2025-08-25T01:55:53.240254Z"
    }
   },
   "cell_type": "code",
   "source": [
    "@torch.no_grad()\n",
    "def encode_texts(texts):\n",
    "    if isinstance(texts, str):\n",
    "        texts = [texts]\n",
    "    tokens = clip_tokenizer(texts, padding=True, return_tensors=\"pt\").to(device)\n",
    "    text_features = clip_model.get_text_features(**tokens) # [B, 512]\n",
    "    text_features = text_features / text_features.norm(p=2, dim=-1, keepdim=True)\n",
    "    return text_features"
   ],
   "id": "fe0350c0e70e1b71",
   "outputs": [],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-25T01:56:29.667634Z",
     "start_time": "2025-08-25T01:55:53.268950Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model = ZhouyuUNet(T=T, time_emb_dim=64)\n",
    "optimizer = AdamW(model.parameters(), lr=1e-4)\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "clip_model.to(device)\n",
    "model.to(device)\n",
    "betas = betas.to(device)\n",
    "alphas = alphas.to(device)\n",
    "alphas_bar = alphas_bar.to(device)\n",
    "sqrt_alphas_bar = sqrt_alphas_bar.to(device)\n",
    "sqrt_one_minus_alphas_bar = sqrt_one_minus_alphas_bar.to(device)\n",
    "\n",
    "@torch.no_grad()\n",
    "def sample(model, text_prompt, num_samples=8):\n",
    "    model.eval()\n",
    "    text_emb = encode_texts(text_prompt)  # [1, 512]\n",
    "    text_emb = text_emb.repeat(num_samples, 1)  # [8, 512]\n",
    "\n",
    "    x = torch.randn(num_samples, 1, 28, 28).to(device)\n",
    "    for t in reversed(range(T)):\n",
    "        t_tensor = torch.full((num_samples,), t, device=device, dtype=torch.long)\n",
    "\n",
    "        noise_pred = model(x, t_tensor, text_emb=text_emb)\n",
    "\n",
    "        alpha = alphas[t]\n",
    "        alpha_bar = alphas_bar[t]\n",
    "        beta = betas[t]\n",
    "        z = torch.randn_like(x) if t > 0 else 0\n",
    "\n",
    "        x = (1 / torch.sqrt(alpha)) * (x - ((1 - alpha) / torch.sqrt(1 - alpha_bar)) * noise_pred) + torch.sqrt(beta) * z\n",
    "\n",
    "    return x.cpu()\n",
    "\n",
    "# 训练\n",
    "os.makedirs(\"samples\", exist_ok=True)\n",
    "epochs = 300\n",
    "print(\"开始训练...\")\n",
    "for epoch in range(epochs):\n",
    "    model.train()\n",
    "    total_loss = 0.0\n",
    "    for step, (images, labels) in enumerate(dataloader):\n",
    "        batch_size = images.shape[0]\n",
    "        images = images.to(device)\n",
    "\n",
    "        # 构造文本描述\n",
    "        texts = [NUM_TO_TEXT[label.item()] for label in labels]\n",
    "        text_emb = encode_texts(texts).detach()  # [B, 512]\n",
    "\n",
    "        # 随机时间步\n",
    "        t = torch.randint(0, T, (batch_size,), device=device)\n",
    "\n",
    "        # 加噪声\n",
    "        sqrt_alpha_bar_t = sqrt_alphas_bar[t].view(-1, 1, 1, 1)\n",
    "        sqrt_one_minus_alpha_bar_t = sqrt_one_minus_alphas_bar[t].view(-1, 1, 1, 1)\n",
    "        noise = torch.randn_like(images)\n",
    "        xt = sqrt_alpha_bar_t * images + sqrt_one_minus_alpha_bar_t * noise\n",
    "\n",
    "        # 预测噪声（传入文本嵌入）\n",
    "        noise_pred = model(xt, t, text_emb=text_emb)\n",
    "        loss = criterion(noise_pred, noise)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "        optimizer.step()\n",
    "\n",
    "        total_loss += loss.item()\n",
    "\n",
    "    avg_loss = total_loss / len(dataloader)\n",
    "    print(f\"Epoch [{epoch+1}/{epochs}], Loss: {avg_loss:.4f}\")\n",
    "\n",
    "    # 每 5 个 epoch 采样一次\n",
    "    if (epoch + 1) % 5 == 0:\n",
    "        gen_images = sample(model, \"手写数字7\", num_samples=8)\n",
    "        gen_images = torch.clamp(gen_images, 0, 1)\n",
    "        grid = torch.cat([img.squeeze() for img in gen_images], dim=1)\n",
    "        img = Image.fromarray((grid.numpy() * 255).astype(np.uint8))\n",
    "        img.save(f\"samples/epoch_{epoch+1:03d}.png\")\n",
    "        print(f\"Sample saved to samples/epoch_{epoch+1:03d}.png\")"
   ],
   "id": "c4bf2f35a05e79f0",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "开始训练...\n",
      "Epoch [1/300], Loss: 0.5904\n",
      "Epoch [2/300], Loss: 0.1637\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[7], line 66\u001B[0m\n\u001B[1;32m     63\u001B[0m loss \u001B[38;5;241m=\u001B[39m criterion(noise_pred, noise)\n\u001B[1;32m     65\u001B[0m optimizer\u001B[38;5;241m.\u001B[39mzero_grad()\n\u001B[0;32m---> 66\u001B[0m \u001B[43mloss\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mbackward\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     67\u001B[0m torch\u001B[38;5;241m.\u001B[39mnn\u001B[38;5;241m.\u001B[39mutils\u001B[38;5;241m.\u001B[39mclip_grad_norm_(model\u001B[38;5;241m.\u001B[39mparameters(), \u001B[38;5;241m1.0\u001B[39m)\n\u001B[1;32m     68\u001B[0m optimizer\u001B[38;5;241m.\u001B[39mstep()\n",
      "File \u001B[0;32m~/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/torch/_tensor.py:648\u001B[0m, in \u001B[0;36mTensor.backward\u001B[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001B[0m\n\u001B[1;32m    638\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m has_torch_function_unary(\u001B[38;5;28mself\u001B[39m):\n\u001B[1;32m    639\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m handle_torch_function(\n\u001B[1;32m    640\u001B[0m         Tensor\u001B[38;5;241m.\u001B[39mbackward,\n\u001B[1;32m    641\u001B[0m         (\u001B[38;5;28mself\u001B[39m,),\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m    646\u001B[0m         inputs\u001B[38;5;241m=\u001B[39minputs,\n\u001B[1;32m    647\u001B[0m     )\n\u001B[0;32m--> 648\u001B[0m \u001B[43mtorch\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mautograd\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mbackward\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    649\u001B[0m \u001B[43m    \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mgradient\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mretain_graph\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcreate_graph\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43minputs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43minputs\u001B[49m\n\u001B[1;32m    650\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/torch/autograd/__init__.py:353\u001B[0m, in \u001B[0;36mbackward\u001B[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001B[0m\n\u001B[1;32m    348\u001B[0m     retain_graph \u001B[38;5;241m=\u001B[39m create_graph\n\u001B[1;32m    350\u001B[0m \u001B[38;5;66;03m# The reason we repeat the same comment below is that\u001B[39;00m\n\u001B[1;32m    351\u001B[0m \u001B[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001B[39;00m\n\u001B[1;32m    352\u001B[0m \u001B[38;5;66;03m# calls in the traceback and some print out the last line\u001B[39;00m\n\u001B[0;32m--> 353\u001B[0m \u001B[43m_engine_run_backward\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    354\u001B[0m \u001B[43m    \u001B[49m\u001B[43mtensors\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    355\u001B[0m \u001B[43m    \u001B[49m\u001B[43mgrad_tensors_\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    356\u001B[0m \u001B[43m    \u001B[49m\u001B[43mretain_graph\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    357\u001B[0m \u001B[43m    \u001B[49m\u001B[43mcreate_graph\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    358\u001B[0m \u001B[43m    \u001B[49m\u001B[43minputs\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    359\u001B[0m \u001B[43m    \u001B[49m\u001B[43mallow_unreachable\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43;01mTrue\u001B[39;49;00m\u001B[43m,\u001B[49m\n\u001B[1;32m    360\u001B[0m \u001B[43m    \u001B[49m\u001B[43maccumulate_grad\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43;01mTrue\u001B[39;49;00m\u001B[43m,\u001B[49m\n\u001B[1;32m    361\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/torch/autograd/graph.py:824\u001B[0m, in \u001B[0;36m_engine_run_backward\u001B[0;34m(t_outputs, *args, **kwargs)\u001B[0m\n\u001B[1;32m    822\u001B[0m     unregister_hooks \u001B[38;5;241m=\u001B[39m _register_logging_hooks_on_whole_graph(t_outputs)\n\u001B[1;32m    823\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m--> 824\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mVariable\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_execution_engine\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mrun_backward\u001B[49m\u001B[43m(\u001B[49m\u001B[43m  \u001B[49m\u001B[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001B[39;49;00m\n\u001B[1;32m    825\u001B[0m \u001B[43m        \u001B[49m\u001B[43mt_outputs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\n\u001B[1;32m    826\u001B[0m \u001B[43m    \u001B[49m\u001B[43m)\u001B[49m  \u001B[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001B[39;00m\n\u001B[1;32m    827\u001B[0m \u001B[38;5;28;01mfinally\u001B[39;00m:\n\u001B[1;32m    828\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m attach_logging_hooks:\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "execution_count": 7
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "gen_images = sample(model, \"手写数字7\", num_samples=8)\n",
    "gen_images = torch.clamp(gen_images, 0, 1) # 将预测结果的像素值截取到[0, 1]\n",
    "gen_images = gen_images * 255\n",
    "\n",
    "fig, axes = plt.subplots(2, 4, figsize=(12, 6))\n",
    "axes = axes.flatten()\n",
    "\n",
    "for i in range(8):\n",
    "    img = gen_images[i].squeeze().cpu().numpy()\n",
    "    axes[i].imshow(img, cmap='gray')\n",
    "    axes[i].axis(\"off\")\n",
    "\n",
    "plt.show()"
   ],
   "id": "df75ee9087cd7a0e"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
